from flask import Blueprint, request, send_file
from PIL import Image
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
import torch
import numpy as np
import glob
import scipy.misc as sc
import imageio
from skimage.transform import resize
from openpyxl import Workbook
from openpyxl.drawing.image import Image as openpyxlImage
from PIL import Image as PilImage

from db import db
from controller.LoginController import getNowUserInfo
from controller.DataController import uploadToOSS
from UNet.configs.config_setting import setting_config as config
from UNet.models.UltraLight_VM_UNet import UltraLight_VM_UNet
from entity.Response import success,error 
from entity.SegmentationHistory import SegmentationHistory
from entity.Log import Log



modelController = Blueprint('modelController',__name__)

FILE_UPLOAD = './controller/LocalImage/'
RESIZE_UPLOAD = './controller/ResizeImage/'
RESULT_PATH = './controller/LocalResult/'
RESUME_PATH = './UNet/results/best.pth'

# 加载模型
model_cfg = config.model_config    
model = UltraLight_VM_UNet(num_classes=model_cfg['num_classes'], 
                           input_channels=model_cfg['input_channels'], 
                           c_list=model_cfg['c_list'], 
                           split_att=model_cfg['split_att'], 
                           bridge=model_cfg['bridge'],)
gpu_ids = [0]
model = torch.nn.DataParallel(model.cuda(), device_ids=gpu_ids, output_device=gpu_ids[0])
best_weight = torch.load(RESUME_PATH, map_location=torch.device('cpu'))
model.module.load_state_dict(best_weight)
print('modle load successfully')

# Parameters
height = 256
width  = 256
channels = 3

"""
上传图片
"""
@modelController.route('/uploadImage',methods=['POST'])
def uploadImage():
    img = request.files['file']
    save_path = FILE_UPLOAD + img.filename
    img.save(save_path)
    resize_save_path = RESIZE_UPLOAD + img.filename
    resize_image(save_path, resize_save_path, size=(256, 256))

    obj = {
        'filename': img.filename,
        'base64': return_img_stream(resize_save_path)
    }
    return success(message='图像上传成功', data=obj)

"""
批量上传图像
"""
@modelController.route('/uploadBatchImage',methods=['POST'])
def uploadBatchImage():
    imgFileList = request.files.getlist('files')
    print(imgFileList)
    dataList = []
    # 保存每一张图片
    for img in imgFileList:
        save_path = FILE_UPLOAD + img.filename
        img.save(save_path)
        resize_save_path = RESIZE_UPLOAD + img.filename
        resize_image(save_path, resize_save_path, size=(256, 256))
        obj = {
            'filename': img.filename,
            'base64': return_img_stream(resize_save_path)
        }
        dataList.append(obj)

    return success(message='图像批量上传成功',data=dataList)

"""
单次图像分割
"""
@modelController.route('/segmentation',methods=['POST'])
def segmentation():
    filename = request.get_json().get('filename')
    image_path = FILE_UPLOAD + filename
    img = imageio.v2.imread(image_path)
    img_resized = resize(img, (height, width, channels), mode='reflect', anti_aliasing=True)

    img_data = np.zeros([1, height, width, channels])
    img_data[0, :, :, :] = img_resized

    img_data = dataset_normalized(img_data)
    img_data = img_data[0]
    img_data = torch.tensor(img_data.copy())
    img_data = img_data.permute( 2, 0, 1)

    img_data = img_data.unsqueeze(0)
    img_data = img_data.cuda(non_blocking=True).float()
    
    out = model(img_data)
    if type(out) is tuple:
        out = out[0]
    out = out.squeeze(1).cpu().detach().numpy()
    
    save_path = RESULT_PATH + 'pred_' + filename
    save_img(out, save_path)
    
    obj = {
        'areaRatio': calAreaRatio(save_path),
        'base64': return_img_stream(save_path)
    }

    # 添加一条分割历史记录 以及一条系统日志
    try:
        resize_path = RESIZE_UPLOAD + filename
        nowUserInfo = getNowUserInfo()
        segmentationHistory = SegmentationHistory(nowUserInfo['username'],uploadToOSS(resize_path),uploadToOSS(save_path),calAreaRatio(save_path))
        db.session.add(segmentationHistory)
        db.session.commit()

        log = Log(nowUserInfo['username'],'进行分割操作')
        db.session.add(log)
        db.session.commit()
    except:
        return success(message='单次分割成功',data=obj)
    
    return success(message='单次分割成功',data=obj)

"""
导出分割结果表
"""
@modelController.route('/exportResult',methods=['POST'])
def exportResult():
    wb = Workbook()
    ws = wb.active

    index = 1
    filenameList = request.get_json().get('filenameList')
    areaRatioList = request.get_json().get('areaRatioList')
    for filename,areaRatio in zip(filenameList,areaRatioList):
        # 设置单元格的宽度和高度（单位：像素）
        ws.column_dimensions['A'].width = 20
        ws.column_dimensions['B'].width = 20
        ws.column_dimensions['C'].width = 20

        ws.row_dimensions[index].height = 80 # 设置行1的高度

        # 加载图片并调整大小以匹配单元格大小
        raw_file_path = FILE_UPLOAD + filename
        raw_img = PilImage.open(raw_file_path)
        raw_img_resized = raw_img.resize((int(ws.column_dimensions['A'].width * 7.5), int(ws.row_dimensions[index].height)))
        resized_img_path = RESULT_PATH + "resize_" + filename
        raw_img_resized.save(resized_img_path)
        raw_img_for_excel = openpyxlImage(resized_img_path)
        ws.add_image(raw_img_for_excel, 'A' + str(index))

        result_file_path = RESULT_PATH + "pred_" + filename
        result_img = PilImage.open(result_file_path)
        result_img_resized = result_img.resize((int(ws.column_dimensions['B'].width * 7.5), int(ws.row_dimensions[index].height)))
        resized_img_path = RESULT_PATH + "resize_pred_" + filename
        result_img_resized.save(resized_img_path)
        result_img_for_excel = openpyxlImage(resized_img_path)
        ws.add_image(result_img_for_excel, 'B' + str(index))

        # 将 areaRatio 写入到 C 列对应的行
        ws['C' + str(index)] = '石细胞面积占比：' + str(areaRatio) + '%'

        index += 1
    
    file_save_path = RESULT_PATH + '分割结果.xlsx'
    wb.save(file_save_path)

    return success(message='生成分割结果表成功')

@modelController.route('/getExportResult',methods=['GET'])
def getExportResult():
    excel_path = RESULT_PATH + '分割结果.xlsx'
    return send_file(excel_path, as_attachment=True, download_name='分割结果.xlsx')

def return_img_stream(img_local_path):
    """
        工具函数:
        获取本地图片流
        :param img_local_path:文件单张图片的本地绝对路径
        :return: 图片流
        """
    import base64
    img_stream = ''
    with open(img_local_path, 'rb') as img_f:
        img_stream = img_f.read()
        img_stream = base64.b64encode(img_stream).decode()
    return img_stream

def save_img(msk_pred, save_path, threshold=0.5):
    msk_pred = np.where(np.squeeze(msk_pred, axis=0) > threshold, 1, 0) 
    # 使用 matplotlib 保存为灰度图像
    plt.imsave(save_path, msk_pred, cmap='gray')
    plt.close()
    
"""
调整图片大小
"""
def resize_image(image_path, output_path, size=(256, 256)):
    """
    将图片调整为指定大小
    """
    with Image.open(image_path) as img:
        img = img.resize(size, Image.ANTIALIAS)  # 使用高质量的缩放算法
        img.save(output_path)

"""
计算所占面积
"""
def calAreaRatio(image_path):
    import cv2
    import numpy as np

    image = cv2.imread(image_path)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    thresh = cv2.threshold(gray,0,255,cv2.THRESH_OTSU + cv2.THRESH_BINARY)[1]
    pixels = cv2.countNonZero(thresh)
    # pixels = len(np.column_stack(np.where(thresh > 0)))

    image_area = image.shape[0] * image.shape[1]
    area_ratio = (pixels / image_area) * 100
    area_ratio = np.round(area_ratio,3)

    print('pixels', pixels)
    print('area ratio', area_ratio)
    return area_ratio

"""
数据normalized
"""
# ===== normalize over the dataset 
def dataset_normalized(imgs):
    imgs_normalized = np.empty(imgs.shape)
    imgs_std = np.std(imgs)
    imgs_mean = np.mean(imgs)
    imgs_normalized = (imgs-imgs_mean)/imgs_std
    for i in range(imgs.shape[0]):
        imgs_normalized[i] = ((imgs_normalized[i] - np.min(imgs_normalized[i])) / (np.max(imgs_normalized[i])-np.min(imgs_normalized[i])))*255
    return imgs_normalized